import torch
import torch.nn as nn
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from sklearn.preprocessing import MinMaxScaler
from LSTM_one.LSTM.dataset_create import create_dataset
from LSTM_one.LSTM.lstm import LSTM

plt.rcParams['font.sans-serif']=['SimHei']  #图中字体改为黑体以兼容中文
plt.rcParams['axes.unicode_minus']=False #负号显示的问题
# 读取数据
data = pd.read_excel('../data/新能源汽车销量.xlsx')
data = data[['销量','市场渗透率']]
# 归一化处理
scaler = MinMaxScaler(feature_range=(0, 1))
scaled_data = scaler.fit_transform(data)

# 定义模型参数
input_size = 2
hidden_size = 120
num_layers = 3
output_size = 2
learning_rate = 0.001
num_epochs = 800

# 定义 LSTM_one 模型
lstm = LSTM(input_size, hidden_size, num_layers, output_size)

look_back = 12
# 加载保存的模型参数
lstm.load_state_dict(torch.load('../model/lstm_model.pth'))

# 将最后的12期数据作为初始输入
test_input = scaled_data[-look_back:]
predictions = []
# 预测未来六期销量，并更新下一期预测数据的输入
for i in range(6):
    # 将数据转换为张量并添加两个新维度（batch 和 feature）
    test_input = torch.from_numpy(test_input).type(torch.Tensor)
    test_input = test_input.unsqueeze(0)  # 在第一个维度和最后一个维度增加维度1
    # 模型预测
    lstm.eval()  # 设置模型为评估模式
    with torch.no_grad():
        prediction = lstm(test_input)  # 注意：使用新的维度
    prediction = prediction.numpy()[0]  # 转换成NumPy数组并取第一个值
    predictions.append(prediction)
    # 更新下一期预测数据的输入
    test_input = np.concatenate([test_input.squeeze().detach().numpy()[1:], [prediction]])  # 去掉新增维度，并在末尾添加预测值
# 反归一化，得到最终的预测结果
predictions = np.array(predictions)
predictions = scaler.inverse_transform(predictions)
#数据处理
for i in range(len(predictions)):
    predictions[i][0] = round(predictions[i][0],0)
    print('未来第',i+1,'月销售预测值:')
    print(predictions[i][0])
    predictions[i][1] = round(predictions[i][1],4)
    print('未来第',i+1,'期市场渗透率:')
    print(predictions[i][1])
#绘制折线图